#include "kernel_operator.h"
using namespace AscendC;
const float FLT_MAX = 1e40;
class KernelThreeNN
{
public:
    __aicore__ inline KernelThreeNN() {}
    __aicore__ inline void Init(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices,
                                uint32_t b,uint32_t n,uint32_t m, uint32_t blockLength,uint32_t CompareLength)
    {
        xyz1Gm.SetGlobalBuffer((__gm__ DTYPE_XYZ1 *)xyz1, blockLength);
        xyz2Gm.SetGlobalBuffer((__gm__ DTYPE_XYZ2 *)xyz2, CompareLength);
        distGm.SetGlobalBuffer((__gm__ DTYPE_DIST *)dist, blockLength);
        indicesGm.SetGlobalBuffer((__gm__ DTYPE_INDICES *)indices, blockLength);
        for (int i = 0;i<b;i++){
            for (int j = 0;j<n;j++){
                float x1 = xyz1Gm.GetValue(i*n*3+j*3);//xyz1[i][j][0]
                float y1 = xyz1Gm.GetValue(i*n*3+j*3+1);//xyz1[i][j][1]
                float z1 = xyz1Gm.GetValue(i*n*3+j*3+2);//xyz1[i][j][2]
                float best1 = FLT_MAX, best2 = FLT_MAX, best3 = FLT_MAX;
                int besti1 = 0, besti2 = 0, besti3 = 0;
                for (int k=0;k<m;k++){
                    float x2 = xyz2Gm.GetValue(i*m*3+k*3);//xyz2[i][k][0];
                    float y2 = xyz2Gm.GetValue(i*m*3+k*3+1);//xyz2[i][k][1];
                    float z2 = xyz2Gm.GetValue(i*m*3+k*3+2);//xyz2[i][k][2];
                    float d = (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1);
                    if (d < best1) {
                        best3 = best2;
                        besti3 = besti2;
                        best2 = best1;
                        besti2 = besti1;
                        best1 = d;
                        besti1 = k;
                    } else if (d < best2) {
                        best3 = best2;
                        besti3 = besti2;
                        best2 = d;
                        besti2 = k;
                    } else if (d < best3) {
                        best3 = d;
                        besti3 = k;
                    }
                }
                distGm.SetValue(i*n*3 + j*3,best1);//dist[i][j][0] = best1;
                indicesGm.SetValue(i*n*3 + j*3,besti1);//idx[i][j][0] = besti1;
                distGm.SetValue(i*n*3 + j*3 + 1,best2);//dist[i][j][1] = best2;
                indicesGm.SetValue(i*n*3 + j*3 + 1,besti2);//idx[i][j][1] = besti2;
                distGm.SetValue(i*n*3 + j*3 + 2,best3);//dist[i][j][2] = best3;
                indicesGm.SetValue(i*n*3 + j*3 + 2,besti3);//idx[i][j][2] = besti3;
            }
        }

    }
private:
    GlobalTensor<DTYPE_XYZ1> xyz1Gm;
    GlobalTensor<DTYPE_XYZ2> xyz2Gm;
    GlobalTensor<DTYPE_DIST> distGm;
    GlobalTensor<DTYPE_INDICES> indicesGm;
};



extern "C" __global__ __aicore__ void three_nn(GM_ADDR xyz1, GM_ADDR xyz2, GM_ADDR dist, GM_ADDR indices, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelThreeNN op;
    op.Init(xyz1, xyz2, dist, indices, tiling_data.b, tiling_data.n, tiling_data.m, tiling_data.blockLength,tiling_data.CompareLength);
    // op.Process();
}



